#include "rclcpp/rclcpp.hpp"
#include "image_service/srv/print_image_path.hpp"
#include <opencv2/opencv.hpp>
#include "task/facenet.h"
#include "task/retinaface.h"

#include <memory>
#include <string>

std::string getFileName(const std::string& path) {
    size_t pos = path.find_last_of("/\\");
    return (pos == std::string::npos) ? path : path.substr(pos + 1);
}

std::string getFileNameWithoutExtension(const std::string& path) {
    std::string fileName = getFileName(path);  // 使用上面的 getFileName 函数
    size_t lastDotPos = fileName.find_last_of('.');
    return (lastDotPos == std::string::npos) ? fileName : fileName.substr(0, lastDotPos);
}

class ImageServiceNode : public rclcpp::Node
{
public:
  ImageServiceNode() : Node("image_service_node")
  {
    service_ = create_service<image_service::srv::PrintImagePath>(
      "print_image_path",
      std::bind(&ImageServiceNode::handle_service, this, std::placeholders::_1, std::placeholders::_2));
  }

private:
    void handle_service(
        const std::shared_ptr<image_service::srv::PrintImagePath::Request> request,
        std::shared_ptr<image_service::srv::PrintImagePath::Response> response)
    {
        RCLCPP_INFO(this->get_logger(), "Received image path: %s", request->image_path.c_str());
        cv::Mat image = cv::imread(request->image_path.c_str());
        std::string fileNameWithoutExt = getFileNameWithoutExtension(request->image_path.c_str());

        if(image.empty())
        {
            response->success = false;
            return;
        }

        /*加载模型权重*/
        retinaface_ = std::make_unique<RetinafaceCustom>();
        retinaface_->LoadModel("/mnt/sdcard/yp/face_recognition_demo/data/retinaFace_mobilenet640_fp.rknn");
        facenet_ = std::make_unique<facenetCustom>();
        facenet_->LoadModel("/mnt/sdcard/yp/face_recognition_demo/data/mobilefacenet.rknn");

        /*模型推理*/
        retinaface_result result;
        retinaface_->Run(image, &result);

        for (int i = 0; i < result.count; ++i) {
            int x1 = result.object[i].box.left;
            int y1 = result.object[i].box.top;
            int x2 = result.object[i].box.right;
            int y2 = result.object[i].box.bottom;
            cv::Rect bbox(x1, y1, x2-x1, y2-y1);
            cv::Mat crop_img = image(bbox);
            cv::Mat align_img;
            retinaface_->FaceAlignment(crop_img, align_img, &result.object[i]);
            float res_vector[128];
            facenet_->Run(align_img, res_vector);
            std::string outPath = "/mnt/sdcard/yp/ros2/ros_service_ws/src/image_service/datasets/" + fileNameWithoutExt + ".dat";
            facenet_->SaveEmbeddings(res_vector, outPath);
        }
        response->success = true;
        RCLCPP_INFO(this->get_logger(), "Processing image");

    }

    rclcpp::Service<image_service::srv::PrintImagePath>::SharedPtr service_;
    std::unique_ptr<RetinafaceCustom> retinaface_;
    std::unique_ptr<facenetCustom> facenet_;
};

int main(int argc, char * argv[])
{
  rclcpp::init(argc, argv);
  rclcpp::spin(std::make_shared<ImageServiceNode>());
  rclcpp::shutdown();
  return 0;
}
